import pygame
import sys
import os
from pygame.locals import *

# 确保chess模块路径正确
current_path = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_path)

from chess.board import ChessBoard
from chess.ai import ChessAI

# 初始化pygame
pygame.init()
pygame.font.init()

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (200, 200, 200)
LIGHT_SQUARE = (240, 217, 181)
DARK_SQUARE = (181, 136, 99)
HIGHLIGHT = (247, 247, 105)
MOVE_HIGHLIGHT = (106, 168, 79, 150)  # 半透明绿色

# 游戏窗口设置
WINDOW_SIZE = (800, 600)
BOARD_SIZE = 480
SQUARE_SIZE = BOARD_SIZE // 8

# 创建游戏窗口
screen = pygame.display.set_mode(WINDOW_SIZE)
pygame.display.set_caption('国际象棋')

# 全局字体
FONT = pygame.font.SysFont('simhei', 32)
SMALL_FONT = pygame.font.SysFont('simhei', 24)

class ChessGame:
    def __init__(self):
        self.screen = screen
        self.clock = pygame.time.Clock()
        self.board = ChessBoard()
        self.ai = ChessAI(3)  # 设置AI深度为3
        
        self.selected_piece = None
        self.valid_moves = []
        self.game_state = 'menu'  # 'menu', 'game', 'game_over'
        self.game_mode = None  # 'human_vs_human', 'human_vs_ai'
        self.current_player = 'white'
        self.winner = None
        
        # 尝试加载棋子图片
        self.piece_images = {}
        self.load_images()
    
    def load_images(self):
        piece_types = ['pawn', 'rook', 'knight', 'bishop', 'queen', 'king']
        colors = ['white', 'black']
        
        # 如果assets目录不存在，则创建
        if not os.path.exists('assets'):
            os.makedirs('assets')
        
        # 检查是否有图片文件，如果没有则使用文本表示
        self.use_text_representation = True
        
        # 字符表示
        self.piece_chars = {
            'white_pawn': '♙', 'white_rook': '♖', 'white_knight': '♘', 
            'white_bishop': '♗', 'white_queen': '♕', 'white_king': '♔',
            'black_pawn': '♟', 'black_rook': '♜', 'black_knight': '♞', 
            'black_bishop': '♝', 'black_queen': '♛', 'black_king': '♚'
        }
    
    def draw_board(self):
        board_offset_x = (WINDOW_SIZE[0] - BOARD_SIZE) // 2
        board_offset_y = (WINDOW_SIZE[1] - BOARD_SIZE) // 2
        
        # 绘制棋盘
        for row in range(8):
            for col in range(8):
                x = board_offset_x + col * SQUARE_SIZE
                y = board_offset_y + row * SQUARE_SIZE
                color = LIGHT_SQUARE if (row + col) % 2 == 0 else DARK_SQUARE
                pygame.draw.rect(self.screen, color, (x, y, SQUARE_SIZE, SQUARE_SIZE))
                
                # 绘制坐标标签
                if col == 0:
                    label = SMALL_FONT.render(f"{8-row}", True, BLACK if (row + col) % 2 == 0 else WHITE)
                    self.screen.blit(label, (x + 5, y + 5))
                if row == 7:
                    label = SMALL_FONT.render(f"{chr(97+col)}", True, BLACK if (row + col) % 2 == 0 else WHITE)
                    self.screen.blit(label, (x + SQUARE_SIZE - 15, y + SQUARE_SIZE - 20))
        
        # 绘制棋子
        for row in range(8):
            for col in range(8):
                piece = self.board.get_piece(row, col)
                if piece:
                    x = board_offset_x + col * SQUARE_SIZE
                    y = board_offset_y + row * SQUARE_SIZE
                    
                    # 高亮选中的棋子
                    if self.selected_piece and self.selected_piece[0] == row and self.selected_piece[1] == col:
                        pygame.draw.rect(self.screen, HIGHLIGHT, (x, y, SQUARE_SIZE, SQUARE_SIZE))
                    
                    # 使用文本表示棋子
                    piece_char = self.piece_chars[f"{piece.color}_{piece.type}"]
                    char_font = pygame.font.SysFont('segoeuisymbol', 48)
                    char_surface = char_font.render(piece_char, True, BLACK)
                    self.screen.blit(char_surface, (x + SQUARE_SIZE//2 - char_surface.get_width()//2, 
                                                  y + SQUARE_SIZE//2 - char_surface.get_height()//2))
        
        # 绘制有效移动位置
        for move in self.valid_moves:
            x = board_offset_x + move[1] * SQUARE_SIZE
            y = board_offset_y + move[0] * SQUARE_SIZE
            
            # 绘制半透明圆形表示可移动位置
            s = pygame.Surface((SQUARE_SIZE, SQUARE_SIZE), pygame.SRCALPHA)
            pygame.draw.circle(s, MOVE_HIGHLIGHT, (SQUARE_SIZE//2, SQUARE_SIZE//2), SQUARE_SIZE//4)
            self.screen.blit(s, (x, y))
    
    def get_square_from_pos(self, pos):
        board_offset_x = (WINDOW_SIZE[0] - BOARD_SIZE) // 2
        board_offset_y = (WINDOW_SIZE[1] - BOARD_SIZE) // 2
        
        # 检查是否点击了棋盘
        if (pos[0] < board_offset_x or pos[0] >= board_offset_x + BOARD_SIZE or
            pos[1] < board_offset_y or pos[1] >= board_offset_y + BOARD_SIZE):
            return None
        
        col = (pos[0] - board_offset_x) // SQUARE_SIZE
        row = (pos[1] - board_offset_y) // SQUARE_SIZE
        return (row, col)
    
    def draw_menu(self):
        self.screen.fill(GRAY)
        
        # 绘制标题
        title = FONT.render("国际象棋", True, BLACK)
        self.screen.blit(title, (WINDOW_SIZE[0]//2 - title.get_width()//2, 100))
        
        # 绘制模式选择按钮
        human_vs_human_btn = pygame.Rect(WINDOW_SIZE[0]//2 - 150, 250, 300, 50)
        human_vs_ai_btn = pygame.Rect(WINDOW_SIZE[0]//2 - 150, 330, 300, 50)
        
        pygame.draw.rect(self.screen, WHITE, human_vs_human_btn)
        pygame.draw.rect(self.screen, WHITE, human_vs_ai_btn)
        
        pygame.draw.rect(self.screen, BLACK, human_vs_human_btn, 2)
        pygame.draw.rect(self.screen, BLACK, human_vs_ai_btn, 2)
        
        hvh_text = FONT.render("双人对战模式", True, BLACK)
        hva_text = FONT.render("人机对战模式", True, BLACK)
        
        self.screen.blit(hvh_text, (human_vs_human_btn.centerx - hvh_text.get_width()//2, 
                                   human_vs_human_btn.centery - hvh_text.get_height()//2))
        self.screen.blit(hva_text, (human_vs_ai_btn.centerx - hva_text.get_width()//2, 
                                   human_vs_ai_btn.centery - hva_text.get_height()//2))
        
        return human_vs_human_btn, human_vs_ai_btn
    
    def draw_game_info(self):
        # 显示当前玩家
        current_player = "白方回合" if self.current_player == 'white' else "黑方回合"
        player_text = SMALL_FONT.render(current_player, True, BLACK)
        self.screen.blit(player_text, (20, 20))
        
        # 显示游戏模式
        mode_text = "双人模式" if self.game_mode == 'human_vs_human' else "人机模式"
        mode_surface = SMALL_FONT.render(mode_text, True, BLACK)
        self.screen.blit(mode_surface, (WINDOW_SIZE[0] - mode_surface.get_width() - 20, 20))
        
        # 添加返回菜单按钮
        back_btn = pygame.Rect(20, WINDOW_SIZE[1] - 40, 120, 30)
        pygame.draw.rect(self.screen, WHITE, back_btn)
        pygame.draw.rect(self.screen, BLACK, back_btn, 2)
        back_text = SMALL_FONT.render("返回菜单", True, BLACK)
        self.screen.blit(back_text, (back_btn.centerx - back_text.get_width()//2, 
                                   back_btn.centery - back_text.get_height()//2))
        
        return back_btn
    
    def draw_game_over(self):
        # 创建半透明叠加层
        overlay = pygame.Surface(WINDOW_SIZE, pygame.SRCALPHA)
        overlay.fill((0, 0, 0, 150))
        self.screen.blit(overlay, (0, 0))
        
        # 显示获胜者信息
        winner_text = f"{'白方' if self.winner == 'white' else '黑方'}获胜!"
        winner_surface = FONT.render(winner_text, True, WHITE)
        self.screen.blit(winner_surface, (WINDOW_SIZE[0]//2 - winner_surface.get_width()//2, 
                                         WINDOW_SIZE[1]//2 - 50))
        
        # 添加重新开始和返回菜单按钮
        restart_btn = pygame.Rect(WINDOW_SIZE[0]//2 - 150, WINDOW_SIZE[1]//2 + 20, 140, 40)
        menu_btn = pygame.Rect(WINDOW_SIZE[0]//2 + 10, WINDOW_SIZE[1]//2 + 20, 140, 40)
        
        pygame.draw.rect(self.screen, WHITE, restart_btn)
        pygame.draw.rect(self.screen, WHITE, menu_btn)
        pygame.draw.rect(self.screen, BLACK, restart_btn, 2)
        pygame.draw.rect(self.screen, BLACK, menu_btn, 2)
        
        restart_text = SMALL_FONT.render("重新开始", True, BLACK)
        menu_text = SMALL_FONT.render("返回菜单", True, BLACK)
        
        self.screen.blit(restart_text, (restart_btn.centerx - restart_text.get_width()//2, 
                                      restart_btn.centery - restart_text.get_height()//2))
        self.screen.blit(menu_text, (menu_btn.centerx - menu_text.get_width()//2, 
                                    menu_btn.centery - menu_text.get_height()//2))
        
        return restart_btn, menu_btn
    
    def handle_menu_event(self, event, human_vs_human_btn, human_vs_ai_btn):
        if event.type == MOUSEBUTTONDOWN:
            if human_vs_human_btn.collidepoint(event.pos):
                self.game_mode = 'human_vs_human'
                self.start_new_game()
            elif human_vs_ai_btn.collidepoint(event.pos):
                self.game_mode = 'human_vs_ai'
                self.start_new_game()
    
    def handle_game_event(self, event, back_btn):
        if event.type == MOUSEBUTTONDOWN:
            if back_btn.collidepoint(event.pos):
                self.game_state = 'menu'
                return
            
            square = self.get_square_from_pos(event.pos)
            if not square:
                return
                
            row, col = square
            
            # 如果已经选择了棋子
            if self.selected_piece is not None:
                # 尝试移动棋子
                for move in self.valid_moves:
                    if move[0] == row and move[1] == col:
                        src_row, src_col = self.selected_piece
                        self.board.move_piece(src_row, src_col, row, col)
                        
                        # 检查游戏结束
                        if self.board.is_checkmate('black' if self.current_player == 'white' else 'white'):
                            self.winner = self.current_player
                            self.game_state = 'game_over'
                        else:
                            # 切换当前玩家
                            self.current_player = 'black' if self.current_player == 'white' else 'white'
                            
                            # 如果是人机模式且轮到AI
                            if self.game_mode == 'human_vs_ai' and self.current_player == 'black':
                                pygame.display.flip()  # 更新显示
                                
                                # AI走棋
                                ai_move = self.ai.get_best_move(self.board, 'black')
                                if ai_move:
                                    src_row, src_col, dst_row, dst_col = ai_move
                                    self.board.move_piece(src_row, src_col, dst_row, dst_col)
                                    
                                    # 检查游戏结束
                                    if self.board.is_checkmate('white'):
                                        self.winner = 'black'
                                        self.game_state = 'game_over'
                                    else:
                                        self.current_player = 'white'
                            
                        self.selected_piece = None
                        self.valid_moves = []
                        return
                
                # 如果点击位置不是有效移动，取消选择
                self.selected_piece = None
                self.valid_moves = []
            
            # 选择新棋子
            piece = self.board.get_piece(row, col)
            if piece and piece.color == self.current_player:
                self.selected_piece = (row, col)
                self.valid_moves = self.board.get_valid_moves(row, col)
    
    def handle_game_over_event(self, event, restart_btn, menu_btn):
        if event.type == MOUSEBUTTONDOWN:
            if restart_btn.collidepoint(event.pos):
                self.start_new_game()
            elif menu_btn.collidepoint(event.pos):
                self.game_state = 'menu'
    
    def start_new_game(self):
        self.board = ChessBoard()
        self.selected_piece = None
        self.valid_moves = []
        self.current_player = 'white'
        self.game_state = 'game'
        self.winner = None
    
    def run(self):
        running = True
        
        while running:
            for event in pygame.event.get():
                if event.type == QUIT:
                    running = False
                
                if self.game_state == 'menu':
                    human_vs_human_btn, human_vs_ai_btn = self.draw_menu()
                    self.handle_menu_event(event, human_vs_human_btn, human_vs_ai_btn)
                
                elif self.game_state == 'game':
                    back_btn = self.draw_game_info()
                    self.handle_game_event(event, back_btn)
                
                elif self.game_state == 'game_over':
                    restart_btn, menu_btn = self.draw_game_over()
                    self.handle_game_over_event(event, restart_btn, menu_btn)
            
            # 绘制游戏界面
            self.screen.fill(GRAY)
            
            if self.game_state == 'menu':
                self.draw_menu()
            
            elif self.game_state == 'game' or self.game_state == 'game_over':
                self.draw_board()
                back_btn = self.draw_game_info()
                
                if self.game_state == 'game_over':
                    restart_btn, menu_btn = self.draw_game_over()
            
            pygame.display.flip()
            self.clock.tick(60)
        
        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = ChessGame()
    game.run() 